from numpy import array
import shapelib, dbflib
from link import Link, Connector
from node import Node, GenNode, SunkNode, ODNode, InterNode
from graph import Graph
from utils import irandomize

class TestModel(object):
	def __init__(self):
		self.sunknode = SunkNode(100)
		self.link = Link(300,2,next=self.sunknode)
		self.gennode = GenNode(self.link)
		
		self.linkids_ = [1,]
		self.links = {1:self.link}
		self.centerlines = {1:[array((0,0),dtype="float"),array((300,0),dtype="float")]}

	def update(self):
		print self.link.snapshot()
		self.link.requota()
		self.sunknode.update()
		self.gennode.update()
		self.link.update()

	def linkids(self):
		return self.linkids_

class Model(object):
	def __init__(self, shapefiles=None):
		if shapefiles:
			self.load(shapefiles)
		else:
			self.init()
	def init(self):
		self.graph = Graph()
		self.links = {}
		self.nodes = {}
		self.zones = {}
		self.connectors = {}
		
		self.coords = {}
		self.centerlines = {}
		self.lanes = {}
		self.capacities = {}
		self.freespeeds = {}
		self.lengths = {}
		self.bbox = None

	def load(self, shapefiles):
		self.init()

		node_file,link_file,zone_file,connector_file=shapefiles
		#read in nodes
		shp = shapelib.ShapeFile(node_file)
		dbf = dbflib.DBFFile(node_file)
		num = shp.info()[0]

		self.bbox = [shp.info()[2][:2], shp.info()[3][:2]]

		for i in range(num):
			record = dbf.read_record(i)
			obj = shp.read_object(i)

			id = record["NO"]
			id = "n"+str(id)
			self.graph.add_node(id)
			self.coords[id] = array(obj.vertices()[0],dtype="float")
			self.nodes[id] = None

		shp.close()
		dbf.close()

		#read in zones
		shp = shapelib.ShapeFile(zone_file)
		dbf = dbflib.DBFFile(zone_file)
		num = shp.info()[0]

		for i in range(num):
			record = dbf.read_record(i)
			obj = shp.read_object(i)

			id = record["NO"]
			id = "z"+str(id)
			self.graph.add_node(id)
			self.coords[id] = array(obj.vertices()[0],dtype="float")
			self.zones[id] = None

		shp.close()
		dbf.close()
		
		#read in links
		shp = shapelib.ShapeFile(link_file)
		dbf = dbflib.DBFFile(link_file)
		num = shp.info()[0]

		for i in range(num):
			record = dbf.read_record(i)
			obj = shp.read_object(i)

			fromid = "n"+str(record["FROMNODENO"])
			toid = "n"+str(record["TONODENO"])
			id = (fromid, toid)
			self.graph.add_arrow(id)
			self.centerlines[id] = array(obj.vertices()[0],dtype="float")
			self.lengths[id] = record["LENGTH"]*1000.0
			self.lanes[id] = record["NUMLANES"]
			self.capacities[id] = float(record["CAPPRT"])/3600.0
			self.freespeeds[id] = float(record["V0PRT"])/3.6
			self.links[id] = None

		shp.close()
		dbf.close()

		#read in connectors
		shp = shapelib.ShapeFile(connector_file)
		dbf = dbflib.DBFFile(connector_file)
		num = shp.info()[0]

		for i in range(num):
			record = dbf.read_record(i)
			obj = shp.read_object(i)

			nodeid = "n"+str(record["NODENO"])
			zoneid = "z"+str(record["ZONENO"])
			if record["DIRECTION"] == "O":
				id = (zoneid, nodeid)
			else:
				id = (nodeid, zoneid)

			self.graph.add_arrow(id)
			self.centerlines[id] = array(obj.vertices()[0],dtype="float")
			self.connectors[id] = None
		
		shp.close()
		dbf.close()

		self.build()
	
	def build(self):
		for id in self.nodes.keys():
			node = InterNode()
			node.id = id #FIXME
			self.nodes[id] = node

		for id in self.zones.keys():
			zone = ODNode()
			zone.id = id#FIXME
			self.zones[id] = zone
			#self.nodes[id] = self.zones[id]

		for linkid in self.links.keys():
			fromid, toid = linkid
			fromnode = self.nodes[fromid]
			tonode = self.nodes[toid]
			
			link = Link(self.lengths[linkid],self.lanes[linkid],
					tonode,self.freespeeds[linkid],
					self.capacities[linkid])
			link.id = linkid#FIXME
			self.links[linkid] = link
			
			fromnode.add_out_link(linkid, link)
			tonode.add_in_link(linkid, link)

		for connectorid in self.connectors.keys():
			fromid, toid = connectorid
			
			if self.nodes.has_key(fromid):
				fromnode = self.nodes[fromid]
			else:
				fromnode = self.zones[fromid]

			if self.nodes.has_key(toid):
				tonode = self.nodes[toid]
			else:
				tonode = self.zones[toid]
			
			connector = Connector(tonode)
			connector.id = connectorid#FIXME
			self.connectors[connectorid] = connector
			#self.links[connectorid] = connector
			
			fromnode.add_out_link(connectorid, connector)
			tonode.add_in_link(connectorid, connector)

	#def links(self):
	#	return self.graph.links()

	#def nodes(self):
	#	return self.graph.nodes()

	def update(self):
		for element in self.zones.values():
			element.requota()
		for element in self.connectors.values():
			element.requota()
		for element in self.nodes.values():
			element.requota()
		for element in self.links.values():
			element.requota()

		for element in irandomize(self.zones.values()):
			element.update()
		for element in irandomize(self.connectors.values()):
			element.update()
		for element in irandomize(self.nodes.values()):
			element.update()
		for element in irandomize(self.links.values()):
			element.update()
		
	def linkids(self):
		return self.links.keys()

	def connectorids(self):
		return self.connectors.keys()

if __name__ == "__main__":
	#model = TestModel()
	#for i in range(150):
	#	model.update()

	shapefiles = [r"D:\myproject\dynat\ljz\ljz_node.shp",
			r"D:\myproject\dynat\ljz\ljz_link.shp",
			r"D:\myproject\dynat\ljz\ljz_zone_centroid.shp",
			r"D:\myproject\dynat\ljz\ljz_connector.shp"]

	model =Model(shapefiles)
	print model.centerlines.keys()
	model.update()
